from ..builder import PIPELINES
import mmcv
import numpy as np
import cv2
from mmcv import imnormalize

@PIPELINES.register_module
class Concat(object):
    """Concat two image.
    Args:
        template_path: template images path
    """

    def __init__(self, template_path):
        self.template_path = template_path
        self.res_mean = 0.0
        self.res_std = 14.47
    def __call__(self, results):
        template_im_name = self.template_path + results['img_info']['filename'].split('.')[0] + '_t.jpg'
        img_temp = mmcv.imread(template_im_name)
        gray_orin = 1.0 * cv2.cvtColor(results['img'], cv2.COLOR_BGR2GRAY)
        gray_temp = 1.0 * cv2.cvtColor(img_temp, cv2.COLOR_BGR2GRAY)
        res = (gray_orin - np.mean(gray_orin)) - (gray_temp - np.mean(gray_temp))
        results['img'] = imnormalize(results['img'], mean=np.array([123.68, 116.779, 103.939]), std=np.array([58.393, 57.12, 57.375]), to_rgb=True)
        # img_temp = imnormalize(img_temp, mean=np.array([123.68, 116.779, 103.939]), std=np.array([58.393, 57.12, 57.375]), to_rgb=True)
        res = (res - self.res_mean) / self.res_std
        results['img'] = np.concatenate([1.0*results['img'], np.expand_dims(res,axis=2)], axis=2)   # 4_channel
        # results['img'] = np.concatenate([np.expand_dims(res, axis=2), np.expand_dims(res, axis=2),np.expand_dims(res, axis=2)], axis=2)
        return results
    def __repr__(self):
        repr_str = self.__class__.__name__
        repr_str += '(template_path={})'.format(
            self.template_path)
        return repr_str